"""
合同目录抽取
"""
import json
from pathlib import Path
from core.llm_answer import answer_with_llm
from core.clean_data import clean_data
from core.text_splitter import TextSplitter
from core.toc_parse import TocParser
from core.contract import Contract
import rapidfuzz
import asyncio


toc_prompt = '''
你是一个专业的合同审查助手。现合同内容如下：
`{}`

你需要基于专业知识，按照合同内容的排列顺序，抽取出能涵盖合同的重要组成部分，整理成大纲目录，目录层级不超过二级。
要求尽量准确可靠，不要输出与大纲目录无关的其他文字。
'''

async def llm_parse_toc():

    contract_data = json.load(open(Path(__file__).parent.parent / r'data/资产转让合同.json', 'r', encoding='utf-8'))
    all_contents = '\n'.join([paragraph['content'] for paragraph in contract_data['textBlocks']])
    text_spliter = TextSplitter(chunk_size=2500, chunk_overlap=0)
    split_docs = text_spliter.create_documents(str(all_contents).strip())
    for doc in split_docs:
        prompt = toc_prompt.format(doc.content)
        print(prompt, '\n ***************************')
        part_toc = await answer_with_llm(prompt, parse_type='str')
        print(part_toc, '\n ***************************')
        for line in str(part_toc).split('\n'):
            content_dict = {
                paragraph['id']: paragraph['content']
                for paragraph in contract_data['textBlocks']
            }
            matched = rapidfuzz.process.extractOne(query=line, choices=content_dict)
            print(line, ' --- ', matched)


def parse_toc_test():
    toc_parser = TocParser()
    for file in Path(r'D:\myworks\gitlab\contract-review\data').rglob('*.json'):
        data = json.load(open(file, 'r', encoding='utf-8'))
        if 'fileId' not in data:
            continue
        data = clean_data(data)
        contract = Contract(
            id=data['fileId'],
            name=data['documentName'],
            paragraphs=data['textBlocks']
        )
        toc = toc_parser.parse(contract)
        print(file.stem)
        print(json.dumps(toc, ensure_ascii=False, indent=4))


if __name__ == '__main__':
    # res = asyncio.run(llm_parse_toc())
    parse_toc_test()


